import sys
import gym
import numpy as np
import torch
from RL import Net

env = gym.make("CartPole-v1")
model = Net().cuda()
pred_net = model.load_state_dict(state_dict=torch.load('dqn.pt'))

times = []
for l in range(30): # 进行100次实验
    observation = env.reset()
    for i in range(1000):
        env.render()
        max_Q = -sys.maxsize
        choesd_action = 0
        for j in range(2):
          action = np.ones((1,1))*j
          X = torch.tensor(np.concatenate((np.tile(observation.reshape((1,4)),(2,1)),np.tile(action,(2,1))),axis=1),dtype=torch.float32).cuda()
          Y = model(X)
          if Y[0].item()>max_Q:
              max_Q = Y[0].item()
              choesd_action = j
        observation, reward, done, info = env.step(choesd_action)
        if done:
          times.append(i+1)
          if i+1 == 500:
              print("game succeed! you have holden on times", i + 1)
          else:
              print("game over! you have holden on times",i+1)
          break
    env.close()
print(np.mean(times))